import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import openpyxl


class BridgeFrequencyPredictor(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, dropout_prob=0.5):
        super(BridgeFrequencyPredictor, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size*2)
        self.fc3 = nn.Linear(hidden_size*2, hidden_size*4)
        self.fc4 = nn.Linear(hidden_size*4, hidden_size*2)
        self.fc5 = nn.Linear(hidden_size*2, hidden_size)
        self.fc6 = nn.Linear(hidden_size, output_size)

        # 添加Dropout层
        self.dropout = nn.Dropout(p=dropout_prob)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = F.relu(self.fc4(x))
        x = F.relu(self.fc5(x))
        x = F.softplus(self.fc6(x))  # 输出层用激活函数,保证输出结果为正数
        return x


def Suspension_bridge_pred(data):
    # 输入数据：主跨跨径（千米)	车道数	桥面宽度（百米）	主梁类型（钢箱梁）	钢加劲梁/钢桁架	钢混组合梁	混凝土梁	桥塔数量	桥塔高度(百米）
    input_size = 3
    hidden_size = 64
    output_size = 3
    # 实例化模型
    pred_model = BridgeFrequencyPredictor(input_size, hidden_size, output_size)
    # 加载模型参数
    pred_model.load_state_dict(torch.load('NN_model_suspension.pth'))
    # 预测
    input_data = torch.tensor(data[0, [0, 7, 8]], dtype=torch.float32)
    frequency = pred_model(input_data)

    damping_ratio = np.zeros([3])   # 阻尼比为跨径的函数 （跨径：km）
    if data[0, 3] == 1 or data[0, 4] == 1:
        damping_ratio[0] = 0.06375 / data[0, 0] + 0.92
        damping_ratio[1] = 0.02407 / data[0, 0] + 0.79
        damping_ratio[2] = 0.01221 / data[0, 0] + 0.6
    else:
        damping_ratio[0] = 0.20702 / data[0, 0] + 0.44
        damping_ratio[1] = 0.19975 / data[0, 0] + 0.33
        damping_ratio[2] = 0.18429 / data[0, 0] + 0.38

    return frequency, damping_ratio


def Cable_stayed_bridge_pred(data):
    # 输入数据：主跨跨径（千米)	车道数	桥面宽度（百米）	主梁类型（钢箱梁）	钢加劲梁/钢桁架	钢混组合梁	混凝土梁	桥塔数量	桥塔高度(百米）
    input_size = 3
    hidden_size = 128
    output_size = 3
    # 实例化模型
    pred_model = BridgeFrequencyPredictor(input_size, hidden_size, output_size)
    # 加载模型参数
    pred_model.load_state_dict(torch.load('NN_model_Cable_stayed.pth'))
    # 预测
    input_data = torch.tensor(data[0, [0, 7, 8]], dtype=torch.float32)
    frequency = pred_model(input_data)
    damping_ratio = np.zeros([3])  # 阻尼比为跨径的函数 （跨径：km）
    if data[0, 3] == 1 or data[0, 4] == 1:
        damping_ratio[0] = 0.06375 / data[0, 0] + 0.92
        damping_ratio[1] = 0.02407 / data[0, 0] + 0.79
        damping_ratio[2] = 0.01221 / data[0, 0] + 0.6
    else:
        damping_ratio[0] = 0.20702 / data[0, 0] + 0.44
        damping_ratio[1] = 0.19975 / data[0, 0] + 0.33
        damping_ratio[2] = 0.18429 / data[0, 0] + 0.38

    return frequency, damping_ratio


def Arch_bridge_pred(data):
    # 输入数据：
    input_size = 5
    hidden_size = 64
    output_size = 3
    # 实例化模型
    pred_model = BridgeFrequencyPredictor(input_size, hidden_size, output_size)
    # 加载模型参数
    pred_model.load_state_dict(torch.load('NN_model_arch.pth'))
    # 预测
    input_data = torch.tensor(data[0, [0, 7, 8, 9, 10]], dtype=torch.float32)
    frequency = pred_model(input_data)
    damping_ratio = np.zeros([3])  # 阻尼比为跨径的函数 （跨径：km）
    if data[0, 3] == 1 or data[0, 4] == 1:
        damping_ratio[0] = 0.06375 / data[0, 0] + 0.92
        damping_ratio[1] = 0.02407 / data[0, 0] + 0.79
        damping_ratio[2] = 0.01221 / data[0, 0] + 0.6
    else:
        damping_ratio[0] = 0.20702 / data[0, 0] + 0.44
        damping_ratio[1] = 0.19975 / data[0, 0] + 0.33
        damping_ratio[2] = 0.18429 / data[0, 0] + 0.38

    return frequency, damping_ratio


def girder_bridge_pred(data):
    # 输入数据：桥梁类型，跨径（m）
    # ================ 25.03.20 修改 ================
    frequency = np.zeros([3])
    if data[0, 0] == 1:
        frequency[0] = 329.85 * data[0, 1] **(-1.427) + 1.35
    else:
        frequency[0] = 329.85 * data[0, 1] **(-1.427) + 1.15

    frequency[1] = 105.08 * data[0, 1] **(-0.912) + 0.2454
    frequency[2] = 105.32 * data[0, 1] ** (-0.867) + 1.221
    # =================================================
    if data[0, 1]<8:
        frequency[0] = 10
        frequency[1] = 13
        frequency[2] = 16

    # 阻尼比
    damping_ratio = np.zeros([3])
    damping_ratio[0] = 2.58 / data[0, 1] + 3.06
    damping_ratio[1] = 2.58 / data[0, 1] + 1.51
    damping_ratio[2] = 2.58 / data[0, 1] + 1.02

    return frequency, damping_ratio


# 测试代码
if __name__ == '__main__':
    # 悬索桥测试
    file_path = 'test_suspension.xlsx'  # 替换为你的 Excel 文件路径
    workbook = openpyxl.load_workbook(file_path)
    sheet = workbook.active
    # 读取数据到 NumPy 数组，并去掉标题行
    data = []
    for row in sheet.iter_rows(min_row=2, values_only=True):  # 从第2行开始读取，去掉标题行
        if any(row):  # 只要有一个有效值，就认为该行是有效的
            data.append(row)  # 从第二列开始，去掉名称列和类型列
    data = np.array(data, dtype=np.float32)

    fre_suspension, dp_suspension = Suspension_bridge_pred(data)

    # 斜拉桥测试
    file_path = 'test_cable_stayed.xlsx'  # 替换为你的 Excel 文件路径
    workbook = openpyxl.load_workbook(file_path)
    sheet = workbook.active
    # 读取数据到 NumPy 数组，并去掉标题行
    data = []
    for row in sheet.iter_rows(min_row=2, values_only=True):  # 从第2行开始读取，去掉标题行
        if any(row):  # 只要有一个有效值，就认为该行是有效的
            data.append(row)  # 从第二列开始，去掉名称列和类型列
    data = np.array(data, dtype=np.float32)

    fre_cable_stayed, dp_calbe_stayed = Cable_stayed_bridge_pred(data)

    # 拱桥测试
    file_path = 'test_arch.xlsx'  # 替换为你的 Excel 文件路径
    workbook = openpyxl.load_workbook(file_path)
    sheet = workbook.active
    # 读取数据到 NumPy 数组，并去掉标题行
    data = []
    for row in sheet.iter_rows(min_row=2, values_only=True):  # 从第2行开始读取，去掉标题行
        if any(row):  # 只要有一个有效值，就认为该行是有效的
            data.append(row)  # 从第二列开始，去掉名称列和类型列
    data = np.array(data, dtype=np.float32)

    fre_arch, dp_arch = Arch_bridge_pred(data)
    print(fre_arch)
    print(dp_arch)

    # 梁桥测试
    file_path = 'test_girder.xlsx'  # 替换为你的 Excel 文件路径
    workbook = openpyxl.load_workbook(file_path)
    sheet = workbook.active
    # 读取数据到 NumPy 数组，并去掉标题行
    data = []
    for row in sheet.iter_rows(min_row=2, values_only=True):  # 从第2行开始读取，去掉标题行
        if any(row):  # 只要有一个有效值，就认为该行是有效的
            data.append(row)  # 从第二列开始，去掉名称列和类型列
    data = np.array(data, dtype=np.float32)

    fre_girder, dp_girder = girder_bridge_pred(data)




